from paddlenlp import Taskflow
import pandas as pd

# 定义抽取的实体类型
schema = ['人物', '事件', '地点','作品','爱好']

# 初始化信息抽取模型
ie = Taskflow(
    "information_extraction",
    schema=schema,
    truncation=True,
    batch_size=16,
    device_id=0
)

# 读取数据
df = pd.read_csv(r"D:\learn\模训\a-certain-nlp-archive\data.csv")

# 处理缺失值并转换为字符串
df["消息内容"] = df["消息内容"].fillna("").astype(str)

# 批量处理文本（自动按batch_size分批次）
texts = df["消息内容"].tolist()
results = ie(texts)

# 定义结果处理函数
def process_entity(entity_list):
    """保留概率>0.5的实体，返回逗号分隔的字符串"""
    return ", ".join([
        item["text"]
        for item in entity_list
        if item.get("probability", 0) > 0.5
    ]) if entity_list else ""

# 创建新列存储结果
df["抽取结果"] = [
    {entity: process_entity(res.get(entity, []))
     for entity in schema}
    for res in results
]

# 保存到CSV
df.to_csv("processed_data.csv", index=False, encoding='utf_8_sig')